#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <cv_bridge/cv_bridge.h>
#include <image_transport/image_transport.h>
#include <std_msgs/String.h>
#include "color_correction/color_correction.hpp"
#include "color_correction/CalibrateColorCorrection.h"
#include "color_correction/ApplyColorCorrection.h"
#include "mcu_correspond/mcu_data_subscriber.h"
#include "mcu_correspond/mcu_data_sender.h"
#include <opencv2/opencv.hpp>
#include <opencv2/imgcodecs.hpp>
#include <locale>

class ColorCorrectionNode {
private:
    ros::NodeHandle nh_;
    ros::NodeHandle pnh_;
    image_transport::ImageTransport it_;
    
    // MCU通信
    std::unique_ptr<McuDataSubscriber> mcu_subscriber_;
    std::unique_ptr<McuDataSender> mcu_sender_;
    
    // Publishers and Subscribers
    image_transport::Subscriber image_sub_;
    ros::Publisher status_pub_;
    
    // Services
    ros::ServiceServer calibrate_service_;
    ros::ServiceServer apply_service_;
    
    // Color corrector instance
    ColorCorrector color_corrector_;
    
    // Parameters
    std::string config_file_path_;
    std::string input_topic_;
    std::string reference_image_path_;
    std::string calibration_input_path_;
    bool use_topic_for_calibration_;  // true: 从话题获取, false: 从文件路径读取
    
    // State tracking
    bool correction_in_progress_;
    cv::Mat latest_image_;  // 存储从话题接收的最新图像
    
public:
    ColorCorrectionNode() : nh_(), pnh_("~"), it_(nh_), correction_in_progress_(false) {
        // Load parameters
        loadParameters();
        
        // Initialize MCU communication
        mcu_subscriber_ = std::make_unique<McuDataSubscriber>(nh_);
        mcu_sender_ = std::make_unique<McuDataSender>(nh_);
        
        // Initialize publishers and subscribers
        if (use_topic_for_calibration_) {
            image_sub_ = it_.subscribe(input_topic_, 1, &ColorCorrectionNode::imageCallback, this);
            ROS_INFO("已订阅图像话题用于校准: %s", input_topic_.c_str());
        } else {
            ROS_INFO("校准将使用文件路径模式，不订阅图像话题");
        }
        status_pub_ = nh_.advertise<std_msgs::String>("/color_correction/status", 1);
        
        // Initialize services
        calibrate_service_ = nh_.advertiseService("/color_correction/calibrate", 
                                                &ColorCorrectionNode::calibrateService, this);
        apply_service_ = nh_.advertiseService("/color_correction/apply_to_file", 
                                           &ColorCorrectionNode::applyService, this);
        
        // Subscribe to MCU correction signal
        mcu_subscriber_->subscribeDouble("coco", 
            [this](double value) {
                if (value == 1.0 && !correction_in_progress_) {
                    ROS_INFO("收到MCU参数校准信号 (coco=%.1f)", value);
                    startColorCorrection();
                }
            });
        
        // Load existing calibration if available
        if (!config_file_path_.empty()) {
            if (color_corrector_.loadParametersFromJson(config_file_path_)) {
                ROS_INFO("色彩校正参数加载成功: %s", config_file_path_.c_str());
            } else {
                ROS_WARN("色彩校正参数加载失败: %s", config_file_path_.c_str());
            }
        }
        
        ROS_INFO("色彩校正节点初始化完成");
        if (use_topic_for_calibration_) {
            ROS_INFO("订阅话题: %s", image_sub_.getTopic().c_str());
        }
        ROS_INFO("MCU通信已启用 - 监听'coco'信号进行参数校准");
        ROS_INFO("注意: 此节点仅负责参数校准，图像校正请使用processor节点");
    }
    
private:
    void loadParameters() {
        pnh_.param<std::string>("config_file_path", config_file_path_, 
                               "/home/firefly/ws_vision/src/color_correction/config/color_correction_params.json");
        pnh_.param<std::string>("input_topic", input_topic_, "/camera/image_raw");
        pnh_.param<std::string>("reference_image_path", reference_image_path_, 
                               "/home/firefly/ws_vision/src/color_correction/reference.jpg");
        pnh_.param<std::string>("calibration_input_path", calibration_input_path_, 
                               "/tmp/current_image.jpg");
        pnh_.param<bool>("use_topic_for_calibration", use_topic_for_calibration_, true);
        
        ROS_INFO("参数加载完成:");
        ROS_INFO("  配置文件路径: %s", config_file_path_.c_str());
        ROS_INFO("  输入话题: %s", input_topic_.c_str());
        ROS_INFO("  参考图像: %s", reference_image_path_.c_str());
        ROS_INFO("  校准输入图像: %s", calibration_input_path_.c_str());
        ROS_INFO("  校准图像来源: %s", use_topic_for_calibration_ ? "ROS话题" : "文件路径");
    }
    
    void startColorCorrection() {
        correction_in_progress_ = true;
        
        ROS_INFO("开始MCU触发的色彩校正参数校准...");
        
        // 执行自动校准流程 - 使用预设的参考图像和当前图像进行校准
        performAutoCalibration();
    }
    
    void finishColorCorrection(bool success) {
        correction_in_progress_ = false;
        
        if (success) {
            // Send success message to LCD
            mcu_sender_->sendJsonString("LCD", "校准完成");
            ROS_INFO("MCU触发的色彩校正参数校准成功完成");
        } else {
            // Send error message to LCD
            mcu_sender_->sendJsonString("LCD", "校准失败");
            ROS_ERROR("MCU触发的色彩校正参数校准失败");
        }
        
        // Reset coco signal to 0.0
        mcu_subscriber_->publishDoubleData("coco", 0.0);
        
        // Publish status
        std_msgs::String status_msg;
        status_msg.data = success ? "参数校准完成" : "参数校准失败";
        status_pub_.publish(status_msg);
    }
    
    void performAutoCalibration() {
        try {
            ROS_INFO("自动校准使用图像:");
            ROS_INFO("  参考图像: %s", reference_image_path_.c_str());
            
            cv::Mat reference_image;
            cv::Mat input_image;
            
            // 加载参考图像
            reference_image = cv::imread(reference_image_path_);
            if (reference_image.empty()) {
                ROS_ERROR("参考图像加载失败: %s", reference_image_path_.c_str());
                finishColorCorrection(false);
                return;
            }
            
            // 根据配置选择输入图像来源
            if (use_topic_for_calibration_) {
                ROS_INFO("  输入图像: 从ROS话题获取 (%s)", input_topic_.c_str());
                
                if (latest_image_.empty()) {
                    ROS_ERROR("尚未从话题接收到图像，请确保话题 %s 正在发布图像", input_topic_.c_str());
                    finishColorCorrection(false);
                    return;
                }
                
                input_image = latest_image_.clone();
                ROS_INFO("使用最新接收的图像进行校准 (尺寸: %dx%d)", input_image.cols, input_image.rows);
                
            } else {
                ROS_INFO("  输入图像: %s", calibration_input_path_.c_str());
                
                input_image = cv::imread(calibration_input_path_);
                if (input_image.empty()) {
                    ROS_ERROR("输入图像加载失败: %s", calibration_input_path_.c_str());
                    finishColorCorrection(false);
                    return;
                }
            }
            
            // 在两幅图像中寻找色卡
            cv::Mat reference_card = color_corrector_.findColorCard(reference_image);
            cv::Mat input_card = color_corrector_.findColorCard(input_image);
            
            if (reference_card.empty()) {
                ROS_ERROR("参考图像中未找到色卡");
                finishColorCorrection(false);
                return;
            }
            
            if (input_card.empty()) {
                ROS_ERROR("输入图像中未找到色卡");
                finishColorCorrection(false);
                return;
            }
            
            // 计算色彩校正矩阵
            bool success = color_corrector_.calculateColorCorrectionMatrix(reference_card, input_card);
            
            if (!success) {
                ROS_ERROR("色彩校正矩阵计算失败");
                finishColorCorrection(false);
                return;
            }
            
            // 保存参数到JSON文件
            if (color_corrector_.saveParametersToJson(config_file_path_)) {
                ROS_INFO("校准参数已保存到: %s", config_file_path_.c_str());
                finishColorCorrection(true);
            } else {
                ROS_ERROR("校准参数保存失败: %s", config_file_path_.c_str());
                finishColorCorrection(false);
            }
            
        } catch (const std::exception& e) {
            ROS_ERROR("自动校准过程中发生异常: %s", e.what());
            finishColorCorrection(false);
        }
    }
    
    void imageCallback(const sensor_msgs::ImageConstPtr& msg) {
        // 当使用话题模式时，持续更新最新图像
        if (use_topic_for_calibration_) {
            try {
                cv_bridge::CvImagePtr cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
                latest_image_ = cv_ptr->image.clone();
                
                // 仅在调试模式下输出
                ROS_DEBUG("已更新最新图像 (尺寸: %dx%d)", latest_image_.cols, latest_image_.rows);
                
                // 如果正在进行校准，可以选择性地保存当前图像到临时文件
                if (correction_in_progress_) {
                    std::string temp_image_path = "/tmp/current_calibration_image.jpg";
                    cv::imwrite(temp_image_path, latest_image_);
                    ROS_INFO("校准过程中，当前图像已保存到: %s", temp_image_path.c_str());
                }
                
            } catch (cv_bridge::Exception& e) {
                ROS_ERROR("图像转换异常: %s", e.what());
            }
        }
    }
    
    bool calibrateService(color_correction::CalibrateColorCorrection::Request& req,
                         color_correction::CalibrateColorCorrection::Response& res) {
        try {
            ROS_INFO("开始色彩校正校准...");
            ROS_INFO("参考图像: %s", req.reference_image_path.c_str());
            ROS_INFO("输入图像: %s", req.input_image_path.c_str());
            
            // Load images
            cv::Mat reference_image = cv::imread(req.reference_image_path);
            cv::Mat input_image = cv::imread(req.input_image_path);
            
            if (reference_image.empty()) {
                res.success = false;
                res.message = "参考图像加载失败: " + req.reference_image_path;
                ROS_ERROR("%s", res.message.c_str());
                return true;
            }
            
            if (input_image.empty()) {
                res.success = false;
                res.message = "输入图像加载失败: " + req.input_image_path;
                ROS_ERROR("%s", res.message.c_str());
                return true;
            }
            
            // Find color cards in both images
            cv::Mat reference_card = color_corrector_.findColorCard(reference_image);
            cv::Mat input_card = color_corrector_.findColorCard(input_image);
            
            if (reference_card.empty()) {
                res.success = false;
                res.message = "参考图像中未找到色卡";
                ROS_ERROR("%s", res.message.c_str());
                return true;
            }
            
            if (input_card.empty()) {
                res.success = false;
                res.message = "输入图像中未找到色卡";
                ROS_ERROR("%s", res.message.c_str());
                return true;
            }
            
            // Calculate color correction matrix
            bool success = color_corrector_.calculateColorCorrectionMatrix(reference_card, input_card);
            
            if (!success) {
                res.success = false;
                res.message = "色彩校正矩阵计算失败";
                ROS_ERROR("%s", res.message.c_str());
                return true;
            }
            
            // Save parameters if path provided
            std::string save_path = req.config_save_path.empty() ? config_file_path_ : req.config_save_path;
            if (color_corrector_.saveParametersToJson(save_path)) {
                ROS_INFO("校准参数保存成功: %s", save_path.c_str());
                
                // Send success notification to MCU
                mcu_sender_->sendJsonKeyValue("LCD", "calibration ok");
            } else {
                ROS_WARN("校准参数保存失败: %s", save_path.c_str());
            }
            
            res.success = true;
            res.message = "色彩校正校准完成成功";
            ROS_INFO("%s", res.message.c_str());
            
            return true;
            
        } catch (const std::exception& e) {
            res.success = false;
            res.message = "Exception during calibration: " + std::string(e.what());
            ROS_ERROR("%s", res.message.c_str());
            return true;
        }
    }
    
    bool applyService(color_correction::ApplyColorCorrection::Request& req,
                     color_correction::ApplyColorCorrection::Response& res) {
        try {
            ROS_INFO("对图像应用色彩校正: %s", req.input_image_path.c_str());
            
            // Load calibration parameters if provided
            if (!req.config_path.empty()) {
                if (!color_corrector_.loadParametersFromJson(req.config_path)) {
                    res.success = false;
                    res.message = "校准参数加载失败: " + req.config_path;
                    ROS_ERROR("%s", res.message.c_str());
                    return true;
                }
            }
            
            if (!color_corrector_.isCalibrationReady()) {
                res.success = false;
                res.message = "色彩校正器未校准";
                ROS_ERROR("%s", res.message.c_str());
                return true;
            }
            
            // Load input image
            cv::Mat input_image = cv::imread(req.input_image_path);
            if (input_image.empty()) {
                res.success = false;
                res.message = "输入图像加载失败: " + req.input_image_path;
                ROS_ERROR("%s", res.message.c_str());
                return true;
            }
            
            // Apply color correction
            cv::Mat corrected_image = color_corrector_.applyColorCorrection(input_image);
            
            // Save corrected image
            if (!cv::imwrite(req.output_image_path, corrected_image)) {
                res.success = false;
                res.message = "校正图像保存失败: " + req.output_image_path;
                ROS_ERROR("%s", res.message.c_str());
                return true;
            }
            
            res.success = true;
            res.message = "色彩校正应用成功并保存到: " + req.output_image_path;
            ROS_INFO("%s", res.message.c_str());
            
            return true;
            
        } catch (const std::exception& e) {
            res.success = false;
            res.message = "色彩校正过程中发生异常: " + std::string(e.what());
            ROS_ERROR("%s", res.message.c_str());
            return true;
        }
    }
};

int main(int argc, char** argv) {
    std::setlocale(LC_ALL, "zh_CN.UTF-8");
    ros::init(argc, argv, "color_correction_node");
    
    try {
        ColorCorrectionNode node;
        ros::spin();
    } catch (const std::exception& e) {
        ROS_FATAL("色彩校正节点失败: %s", e.what());
        return -1;
    }
    
    return 0;
}
